import os
import zipfile
from pathlib import Path
import json
import tkinter as tk
from tkinter import messagebox, filedialog, colorchooser
import pyautogui
from pynput import keyboard
from PIL import Image, ImageDraw, ImageFont
import threading


SAVE_DIR = "screenshots"
CONFIG_FILE = "config.json"  # 配置文件路径


def load_config(config_file):
    """
    加载配置文件，如果不存在则创建并写入默认配置
    :param config_file: 配置文件路径
    :return: 配置字典
    """
    default_config = {
        "watermark_text": "202472610119 于洪朋",
        "font_path": "simsun.ttc",
        "font_size": 20,
        "text_color": [0, 0, 0, 165],
        "angle": 45
    }

    if not os.path.exists(config_file):
        print(f"配置文件 {config_file} 不存在，创建并写入默认配置。")
        try:
            # 确保保存配置文件的目录存在
            os.makedirs(Path(config_file).parent, exist_ok=True)
            with open(config_file, 'w', encoding='utf-8') as f:
                json.dump(default_config, f, ensure_ascii=False, indent=4)
            print(f"已创建配置文件: {config_file}")
            return default_config
        except Exception as e:
            print(f"创建配置文件失败: {e}")
            raise
    else:
        try:
            with open(config_file, 'r', encoding='utf-8') as f:
                config = json.load(f)
            # 确保颜色是列表形式
            if "text_color" in config and isinstance(config["text_color"], str):
                config["text_color"] = list(eval(config["text_color"]))  # 如果颜色是字符串形式的元组，转换为列表
            return config
        except Exception as e:
            print(f"加载配置文件失败: {e}")
            # 如果加载失败，返回默认配置
            return default_config


def add_watermark(image_path, watermark_text, font_path, font_size, text_color, angle=0):
    """
    添加水印到图片
    :param image_path: 图片路径
    :param watermark_text: 水印文字
    :param font_path: 字体路径
    :param font_size: 字体大小
    :param text_color: 文字颜色 (RGBA)
    :param angle: 旋转角度
    """
    try:
        image = Image.open(image_path).convert("RGBA")
        
        watermark = Image.new("RGBA", image.size, (0, 0, 0, 0))
        
        draw = ImageDraw.Draw(watermark)
        
        try:
            font = ImageFont.truetype(font_path, font_size)
        except IOError:
            print(f"字体文件 {font_path} 未找到，使用默认字体。")
            font = ImageFont.load_default()

        bbox = font.getbbox(watermark_text)
        text_width = bbox[2] - bbox[0]
        text_height = bbox[3] - bbox[1]

        text_layer = Image.new("RGBA", (text_width, text_height), (0, 0, 0, 0))
        text_draw = ImageDraw.Draw(text_layer)
        text_draw.text((0, 0), watermark_text, font=font, fill=text_color)
        rotated_layer = text_layer.rotate(
            angle, 
            expand=True, 
            fillcolor=(0, 0, 0, 0),  
            resample=Image.BILINEAR
        )
        rw, rh = rotated_layer.size 
        step_x = int(rw * 1.5)
        step_y = int(rh * 1.5)
        for x in range(-rw, image.width + rw, step_x):
            for y in range(-rh, image.height + rh, step_y):
                pos = (
                    x - rw // 2,
                    y - rh // 2
                )
                watermark.paste(rotated_layer, pos)
        
        watermarked = Image.alpha_composite(image, watermark)
        if image_path.lower().endswith(('.jpg', '.jpeg')):
            watermarked = watermarked.convert('RGB')
        
        watermarked.save(image_path)
        print(f"水印已添加到 {image_path}")
    except Exception as e:
        print(f"添加水印失败: {e}")


def extract_images_from_docx(docx_path, output_folder, config):
    """
    从Word文档中提取图片到指定文件夹，并添加水印
    :param docx_path: Word文档路径
    :param output_folder: 图片输出文件夹路径
    :param config: 配置字典
    """
    # 确保输出文件夹存在
    Path(output_folder).mkdir(parents=True, exist_ok=True)
    
    try:
        # 将.docx文件作为ZIP打开
        with zipfile.ZipFile(docx_path, 'r') as zip_ref:
            # 遍历ZIP内所有文件
            for file_info in zip_ref.infolist():
                # 检查文件是否在media目录且为图片
                if file_info.filename.startswith('word/media/'):
                    # 获取文件名
                    file_name = os.path.basename(file_info.filename)
                    # 构建输出路径
                    output_path = os.path.join(output_folder, file_name)
                    # 确保不覆盖已有文件（可选：添加编号处理重名）
                    counter = 1
                    base, ext = os.path.splitext(file_name)
                    while os.path.exists(output_path):
                        new_name = f"{base}_{counter}{ext}"
                        output_path = os.path.join(output_folder, new_name)
                        counter += 1
                    # 提取文件
                    with zip_ref.open(file_info) as source, open(output_path, 'wb') as target:
                        target.write(source.read())
                    
                    print(f"提取图片: {file_name} -> {output_path}")
                    # 添加水印
                    add_watermark(
                        output_path,
                        config["watermark_text"],
                        config["font_path"],
                        config["font_size"],
                        tuple(config["text_color"]),
                        config["angle"]
                    )
        print("图片提取完成！")
    except zipfile.BadZipFile:
        print("错误：文件不是有效的ZIP文档，请确认是否为.docx格式")
    except FileNotFoundError:
        print("错误：指定的Word文档不存在")


# 使用示例
if __name__ == "__main__":
    # 加载配置
    config = load_config(CONFIG_FILE)

    # 示例：手动指定docx文件和输出目录
    docx_file = "www.docx"  # 修改为你的Word文件路径
    output_dir = "1"  # 图片输出目录

    if not os.path.exists(docx_file):
        print(f"文件 {docx_file} 不存在，请检查路径。")
    else:
        extract_images_from_docx(docx_file, output_dir, config)